package main

import (
  {% if is_test or to_json %}
	"encoding/json"
	{% endif %}
	"fmt"
	"os"
	"strconv"
)
{% if is_test or to_json %}
type response struct {
	Predict int `json:"predict"`
	PredictProba []float64 `json:"predict_proba"`
}
{% endif %}

func findMax(nums []int) int {
	var idx = 0
	for i := 0; i < len(nums); i++ {
		if nums[i] > nums[idx] {
			idx = i
		}
	}
	return idx
}

func normVals(nums []int) []float64 {
	var l = len(nums)
	var sum = 0
	for i := 0; i < l; i++ {
		sum += nums[i]
	}
	result := make([]float64, l)
	if sum == 0 {
		for i := 0; i < l; i++ {
			result[i] = float64(1) / float64(l)
		}
	} else {
		for i := 0; i < l; i++ {
			result[i] = float64(nums[i]) / float64(sum)
		}
	}
	return result
}

func compute(features []float64) []int {
	classes := make([]int, {{ n_classes }})
	{{ tree | indent(4, True) }}
	return classes
}

func Predict(features []float64) int {
	return findMax(compute(features))
}

func PredictProba(features []float64) []float64 {
	return normVals(compute(features))
}

func main() {

	// Features:
	var features []float64
	for _, arg := range os.Args[1:] {
		if n, err := strconv.ParseFloat(arg, 64); err == nil {
			features = append(features, n)
		}
	}

	{% if is_test or to_json %}
	// Get JSON:
	prediction := Predict(features)
	probabilities := PredictProba(features)
	res, _ := json.Marshal(&response{Predict: prediction, PredictProba: probabilities})
	fmt.Println(string(res))
	{% else %}
	// Get class prediction:
	prediction := Predict(features)
	fmt.Printf("Predicted class: #%d\n", prediction)

	// Get class probabilities:
	probabilities := PredictProba(features)
	for i := 0; i < len(probabilities); i++ {
		fmt.Printf("Probability of class #%d : %.f\n", i, probabilities[i])
	}
	{% endif %}

}